/*
 * Copyright (c) 2025.
 * All rights reserved.
 */

#include "driver.h"
#include "registry_monitor.h"
#include "ioctl_communicator.h"
#include "string_utils.h"

#define MAX_PATH_LENGTH		1024
#define CALLBACK_ALTITUDE	L"321999"

LARGE_INTEGER	g_RegisterCookie = { 0 };
UNICODE_STRING	g_RegistryPath;
UNICODE_STRING	g_UninstallKey = RTL_CONSTANT_STRING(L"UninstallString");

BOOLEAN GetRegistryPath(
	PVOID pRegistryObj)
{
	if (pRegistryObj == NULL || FALSE == MmIsAddressValid(pRegistryObj)) {
		return FALSE;
	}
	// allocate memory
	PVOID lpObjectNameInfo = ExAllocatePool(NonPagedPool, MAX_PATH_LENGTH);
	if (lpObjectNameInfo == NULL) {
		return FALSE;
	}
	// get regiter path
	ULONG ulRetLen = 0;
	NTSTATUS status = ObQueryNameString(pRegistryObj, (POBJECT_NAME_INFORMATION)lpObjectNameInfo, MAX_PATH_LENGTH, &ulRetLen);
	if (!NT_SUCCESS(status)) {
		ExFreePool(lpObjectNameInfo);
		return FALSE;
	}
	// copy path
	RtlCopyUnicodeString(&g_RegistryPath, (PUNICODE_STRING)lpObjectNameInfo);
	// release memory
	ExFreePool(lpObjectNameInfo);
	return TRUE;
}

NTSTATUS InstallRegistryCallback(
	_In_ PDRIVER_OBJECT driverObj)
{
	if (driverObj == NULL) {
		return STATUS_INVALID_PARAMETER;
	}

	NTSTATUS status;
	UNICODE_STRING Altitude = { 0 };

	RtlInitUnicodeString(&Altitude, CALLBACK_ALTITUDE);
	status = CmRegisterCallbackEx(RegistryCallback, &Altitude, driverObj, NULL, &g_RegisterCookie, NULL);
	if (STATUS_SUCCESS != status) {
		DBG_PRINT("CmRegisterCallbackEx: Fail\n");
		return status;
	}

	if (!InitUnicodeString(&g_RegistryPath, MAX_PATH_LENGTH, 'reg')) {
		CmUnRegisterCallback(g_RegisterCookie);
		DBG_PRINT("CmRegisterCallbackEx: Allocate buffer fail!!!\n");
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	return STATUS_SUCCESS;
}
VOID UninstallRegistryCallback()
{
	CmUnRegisterCallback(g_RegisterCookie);
	RtlFreeUnicodeString(&g_RegistryPath);
}

NTSTATUS RegistryCallback(
	_In_ PVOID callbackCtx,
	_In_ PVOID argument1,
	_In_ PVOID argument2)
{
	UNREFERENCED_PARAMETER(callbackCtx);

	if (g_RegistryPath.Buffer == NULL) {
		return STATUS_INVALID_PARAMETER;
	}

	LONG OptType = (REG_NOTIFY_CLASS)(ULONG_PTR)argument1;
	if (OptType != RegNtPostSetValueKey) {
		return STATUS_SUCCESS;
	}
	PREG_POST_OPERATION_INFORMATION PostInfo = (PREG_POST_OPERATION_INFORMATION)argument2;
	if (PostInfo == NULL) {
		return STATUS_SUCCESS;
	}
	PREG_SET_VALUE_KEY_INFORMATION PreSetValueInfo = (PREG_SET_VALUE_KEY_INFORMATION)PostInfo->PreInformation;
	if (PreSetValueInfo == NULL) {
		return STATUS_SUCCESS;
	}
	PUNICODE_STRING valueName = PreSetValueInfo->ValueName;
	if (valueName == NULL) {
		return STATUS_SUCCESS;
	}
	if (RtlCompareUnicodeString(valueName, &g_UninstallKey, TRUE) != 0) {
		return STATUS_SUCCESS;
	}
	if (!GetRegistryPath(PreSetValueInfo->Object)) {
		return STATUS_SUCCESS;
	}
	if (wcsstr(g_RegistryPath.Buffer, L"\\CurrentVersion\\Uninstall\\") != NULL) {
		DBG_PRINT("[Uninstall Set],%wZ,%wZ \n", PreSetValueInfo->ValueName, g_RegistryPath);
		PMSG_ITEM	msgItem;
		MSG_PARAM   msgParam = { MSG_TYPE_INSTALL_END, &g_RegistryPath, NULL, NULL };
		IoctlSendMsg(msgParam, FALSE, &msgItem);
	}

	return STATUS_SUCCESS;
}